/**
 * @file   Spline.h
 * @author xingyifan
 * @date   2021/11/21 14:55:22
 *
 * @brief  spline
 *
 *
 */

#ifndef __SPLINE__
#define __SPLINE__

#include "Config.h"
#include "Interp.h"
#include "Vec.h"

/**
 * @file   Spline.h
 * @author xingyifan
 * @date   2021/11/22 22:20:18
 *
 * @brief  向量及矩阵输出函数，和通过 lapack 求解线性方程组
 *
 *
 */

vector<double> solve(vector<vector<double>> &, vector<double> &);
void vec_out(vector<double> &);
void mat_out(vector<vector<double>> &);
double max_norm(vector<double> &);

/**
 * @file   Spline.h
 * @author xingyifan
 * @date   2021/11/22 22:21:05
 *
 * @brief  Spline
 *
 *
 */

template <int Dim, int Order, SplineType t>
class Spline
{
private:
    using CoefType = Vec<double, Dim>;
    vector<double> interval;                         //非零区间
    vector<Polynomial<Order, CoefType>> spline_poly; //分段多项式

public:
    Spline() {}
    Spline(const vector<Polynomial<Order, CoefType>> &spline_poly) : spline_poly(spline_poly) {}
    template <int Order2, SplineType t2>
    Spline(const Spline<Dim, Order2, t2> &spline)
    {
        interval = spline.get_interval();
        vector<Polynomial<Order2, CoefType>> _spline_poly = spline.get_spline();
        spline_poly.resize(_spline_poly.size());
        for (int i = 0; i < _spline_poly.size(); i++)
        {
            spline_poly[i] = _spline_poly[i];
        }
    }

    //二维样条插值
    template <int Ord>
    friend Spline<2, Ord, ppForm> fitCurve(const vector<Vec<double, 2>> &, BCType);

    //标准样条插值
    template <int Ord, SplineType _t>
    friend Spline<1, Ord, _t> interpolate(const INTERPC &, BCType);

    //基数B样条
    template <int Ord>
    friend Spline<1, Ord, cardinalB> B_spline(int i);

    //计算部分差分表
    vector<vector<double>> calculateTable(int width, INTERPC &interpC_);

    //逐段计算样条，同时通过每段多项式添加新的边值条件
    Polynomial<Order, double> pieceCalculate(INTERPC &interpC, vector<double> &M, int index);

    //导数
    Spline<Dim, Order - 1, t> derivation()
    {
        Spline<Dim, Order - 1, t> res;
        res.resize(spline_poly.size());
        for (int i = 0; i < res.size(); i++)
        {
            res[i] = spline_poly[i].derivation();
        }
        //传递区间
        res.set_interval(interval);
        return res;
    }

    //计算值
    CoefType evaluate(double x)
    {
        if (x == interval[0])
        {
            return spline_poly[0].evaluate(x);
        }
        return get_poly(x).evaluate(x);
    }

    //取分段多项式
    Polynomial<Order, CoefType> operator[](u_int index) const
    {
        if (index < spline_poly.size())
        {
            return spline_poly[index];
        }
        return spline_poly[0];
    }

    //修改多项式
    Polynomial<Order, CoefType> &operator[](u_int index)
    {
        if (index < spline_poly.size())
        {
            return spline_poly[index];
        }
        return spline_poly[0];
    }

    //根据点所在区间取多项式
    Polynomial<Order, CoefType> get_poly(double x) const
    {
        for (int i = 0; i < interval.size() - 1; i++)
        {
            //按左开右闭区间
            if (x > interval[i] && x <= interval[i + 1])
            {
                return spline_poly[i];
            }
        }
        return Polynomial<Order, CoefType>();
    }

// B样条加法：注意未初始化的样条
//注意调整区间和多项式向量的长度
#define SPLINE_BINARY_OP(OpName, Op)                                             \
    template <int Order2>                                                        \
    auto OpName(const Spline<1, Order2, cardinalB> &spline2)                     \
    {                                                                            \
        Spline<1, MAX(Order, Order2), cardinalB> res;                            \
        if (interval.empty())                                                    \
        {                                                                        \
            return Spline<1, MAX(Order, Order2), cardinalB>(spline2);            \
        }                                                                        \
        if (spline2.get_interval().empty())                                      \
        {                                                                        \
            return Spline<1, MAX(Order, Order2), cardinalB>(*this);              \
        }                                                                        \
        int left = MIN(interval.front(), spline2.get_interval().front());        \
        int right = MAX(interval.back(), spline2.get_interval().back());         \
        vector<double> interval;                                                 \
        for (int i = left; i < right + 1; i++)                                   \
        {                                                                        \
            interval.push_back(i);                                               \
        }                                                                        \
        res.set_interval(interval);                                              \
        res.resize(interval.size() - 1);                                         \
        for (int i = 1; i < interval.size(); i++)                                \
        {                                                                        \
            res[i - 1] = get_poly(interval[i]) Op spline2.get_poly(interval[i]); \
        }                                                                        \
        return res;                                                              \
    }

    SPLINE_BINARY_OP(operator+, +)
    SPLINE_BINARY_OP(operator-, -)
#undef SPLINE_BINARY_OP

//样条乘法
#define SPLINE_NUM_OP(OpName, Op)                        \
    Spline<Dim, Order, t> OpName(const CoefType k) const \
    {                                                    \
        Spline<Dim, Order, t> res = *this;               \
        for (int i = 0; i < spline_poly.size(); i++)     \
        {                                                \
            res[i] = res[i] Op k;                        \
        }                                                \
        return res;                                      \
    }

    SPLINE_NUM_OP(operator*, *)
    SPLINE_NUM_OP(operator/, /)
#undef SPLINE_NUM_OP

    //多项式乘样条
    template <int Order2>
    //传参时如果不改变参数值，尽可能把参数写成const，防止出现传入常量时编译产生的问题
    Spline<Dim, Order + Order2, t> operator*(const Polynomial<Order2, CoefType> &poly)
    {
        Spline<Dim, Order + Order2, t> res = *this;
        for (int i = 0; i < spline_poly.size(); i++)
        {
            res[i] = res[i] * poly;
        }
        return res;
    }

    //样条赋值
    template <int Order2, SplineType t2>
    void operator=(const Spline<Dim, Order2, t2> &spline)
    {
        interval = spline.get_interval();
        vector<Polynomial<Order2, CoefType>> _spline_poly = spline.get_spline();
        //清空重新赋值
        spline_poly.resize(_spline_poly.size());
        for (int i = 0; i < _spline_poly.size(); i++)
        {
            spline_poly[i] = _spline_poly[i];
        }
    }

    //右移样条
    Spline<1, Order, t> operator>>(double k)
    {
        Spline<1, Order, t> res = *this;
        for (int i = 0; i < spline_poly.size(); i++)
        {
            res[i] = res[i] >> k;
        }
        vector<double> res_interval = interval;
        for (int i = 0; i < res_interval.size(); i++)
        {
            res_interval[i] += k;
        }
        res.set_interval(res_interval);
        return res;
    }

    //标准输出
    friend ostream &operator<<(ostream &out, const Spline &spline)
    {
        for (int i = 0; i < spline.size(); i++)
        {
            out << "s" << i + 1 << "(x) = " << spline[i];
        }
        return out;
    }

    //绘制图像
    void plot(char color)
    {
        for (int i = 0; i < spline_poly.size(); i++)
        {
            spline_poly[i].plot(interval[i], interval[i + 1], color);
            cout << "hold on;" << endl;
        }
    }

    void plot2(char color)
    {
        for (int i = 0; i < spline_poly.size(); i++)
        {
            spline_poly[i].plot2(interval[i], interval[i + 1], color);
            cout << "hold on;" << endl;
        }
    }

    //获取和修改成员变量
    u_int dim() const { return Dim; }
    u_int order() const { return Order; }
    u_int size() const { return spline_poly.size(); }
    void resize(u_int n) { spline_poly.resize(n); }
    void set_interval(const vector<double> &_interval) { interval = _interval; }
    vector<double> get_interval() const { return interval; }
    vector<Polynomial<Order, CoefType>> get_spline() const { return spline_poly; }
};

/**
 * @date   2021/11/23 20:15:54
 *
 * @arg
 * @return
 *
 */

//计算基数B样条
template <int Ord>
Spline<1, Ord, cardinalB> B_spline(int i)
{
    Spline<1, Ord, cardinalB> res;
    res.set_interval({i - 1.0, i * 1.0});
    res.resize(1);
    res[0] = Polynomial<Ord, Vec1>({1});

    //利用递推公式和平移性质
    for (int j = 0; j < Ord; j++)
    {
        //注意每次阶数 Order = j + 1
        res = res * Polynomial<Ord, Vec1>({(1.0 - i) / (j + 1), 1.0 / (j + 1)}) +
              (res >> 1) * Polynomial<Ord, Vec1>({1.0 * (j + 1 + i) / (j + 1), -1.0 / (j + 1)});
    }
    return res;
}

/**
 * @date   2021/11/23 20:14:32
 *
 * @arg    int width, INTERPC &interpC_
 * @return vector<vector<double>>
 *
 */

template <int Dim, int Order, SplineType t>
vector<vector<double>> Spline<Dim, Order, t>::calculateTable(int width, INTERPC &interpC_)
{
    //存放点作为样条区间
    interval = interpC_.points;
    spline_poly.resize(interval.size() - 1);

    vector<vector<double>> res;
    vector<double> cond;
    for (int i = 0; i < interval.size(); i++)
    {
        for (int j = 0; j < interpC_.interpSites[interval[i]].size(); j++)
        {
            cond.push_back(interval[i]);
        }
    }
    for (int i = 0; i < cond.size(); i++)
    {
        //注意边界条件取最小值
        for (int j = 0; j < MIN(width, i + 1); j++)
        {
            if (j == 0)
            {
                res.push_back({interpC_.interpSites[cond[i]][0]});
                continue;
            }
            double left = cond[i - j], right = cond[i];
            if (left != right)
            {
                res[i].push_back((res[i][j - 1] - res[i - 1][j - 1]) / (right - left));
                continue;
            }
            int factor = 1;
            for (int k = 1; k <= j; factor *= k, k++)
                ;
            res[i].push_back(1.0 / factor * interpC_.interpSites[cond[i]][factor]);
        }
    }
    return res;
}

/**
 * @date   2021/11/23 20:12:49
 *
 * @arg    INTERPC &interpC, vector<double> &M, int index
 * @return Polynomial<Order, double>
 *
 */

template <int Dim, int Order, SplineType t>
Polynomial<Order, double> Spline<Dim, Order, t>::pieceCalculate(INTERPC &interpC, vector<double> &M, int index)
{
    Polynomial<Order, double> res;
    Polynomial<0, double> res_0({interpC.interpSites[interpC.points[index]][0]});
    Polynomial<1, double> res_1({-1 * interpC.points[index], 1});
    Polynomial<2, double> res_2 = res_1.pow<2>();
    Polynomial<3, double> res_3 = res_1.pow<3>();

    //利用泰勒展开式
    // s_i(x) = f_i + f[x_i, x_i](x - x_i) + M_i / 2 * (x - x_i)^2 + (M_{i+1} - M_i) / 6(x_{i+1} - x_i) * (x - x_i)^3
    res = res_0 + res_1 * interpC.interpSites[interpC.points[index]][1] + res_2 * M[index] / 2 +
          res_3 * (M[index + 1] - M[index]) / (6 * (interpC.points[index + 1] - interpC.points[index]));
    //添加新的边值条件
    double m_1 = res.derivation().evaluate(interpC.points[index + 1]);
    interpC.interpSites[interpC.points[index + 1]].push_back(m_1);

    return res;
}

/**
 * @date   2021/11/22 22:21:30
 *
 * @arg    const INTERPC &interpC, BCType type
 * @return Spline<1, Ord, _t>
 *
 */

//标准三类样条插值和二次、三次B样条
template <int Ord, SplineType t>
Spline<1, Ord, t> interpolate(const INTERPC &interpC, BCType type)
{
    Spline<1, Ord, t> res_spline;
    INTERPC interpC_ = interpC;
    const vector<double> points = interpC_.points;

    switch (t)
    {
    case ppForm:
    {
        //插值点的个数
        int N = interpC_.interpSites.size();
        //计算二阶导的关系
        vector<vector<double>> tableOfDividedDiffs_ = res_spline.calculateTable(Ord, interpC_);
        // mat_out(tableOfDividedDiffs_);

        vector<vector<double>> A(N, vector<double>(N));
        vector<double> b(N);
        vector<double> lam(N - 2), mu(N - 2);

        //计算系数
        for (int i = 0; i < N - 2; i++)
        {
            mu[i] = (points[i + 1] - points[i]) / (points[i + 2] - points[i]);
            lam[i] = (points[i + 2] - points[i + 1]) / (points[i + 2] - points[i]);
        }
        //系数矩阵
        for (int i = 1; i < N - 1; i++)
        {
            A[i][i - 1] = mu[i - 1];
            A[i][i + 1] = lam[i - 1];
            A[i][i] = 2;
        }

        //完全样条
        if (type == complete)
        {
            //构造线性方程组
            A[0][0] = 2, A[0][1] = 1;
            A[N - 1][N - 2] = 1, A[N - 1][N - 1] = 2;

            for (int i = 0; i < N; i++)
            {
                b[i] = 6 * tableOfDividedDiffs_[i + 2][Ord - 1];
            }
        }
        //无结样条
        if (type == notAknot)
        {
            //构造线性方程组
            //利用三阶导相等
            A[0][0] = points[2] - points[1], A[0][1] = points[0] - points[2], A[0][2] = points[1] - points[0];
            A[N - 1][N - 3] = points[N - 1] - points[N - 2], A[N - 1][N - 2] = points[N - 3] - points[N - 1], A[N - 1][N - 1] = points[N - 2] - points[N - 3];

            for (int i = 1; i < N - 1; i++)
            {
                b[i] = 6 * tableOfDividedDiffs_[i + 1][Ord - 1];
            }
            b[0] = b[N - 1] = 0;
        }
        //周期样条
        if (type == periodic)
        {
            //构造线性方程组
            A[0][0] = 1, A[0][N - 1] = -1;
            A[N - 1][0] = 2 * (points[1] - points[0]), A[N - 1][1] = points[1] - points[0];
            A[N - 1][N - 2] = points[N - 1] - points[N - 2], A[N - 1][N - 1] = 2 * (points[N - 1] - points[N - 2]);

            for (int i = 1; i < N - 1; i++)
            {
                b[i] = 6 * tableOfDividedDiffs_[i + 1][Ord - 1];
            }
            // 6(f[x_1, x_2] - f[x_N - x_{N-1}])
            b[0] = 0, b[N - 1] = 6 * ((interpC_.interpSites[points[1]][0] - interpC_.interpSites[points[0]][0]) / (points[1] - points[0]) -
                                      (interpC_.interpSites[points[N - 1]][0] - interpC_.interpSites[points[N - 2]][0]) / (points[N - 1] - points[N - 2]));
        }

        // vec_out(b);
        // cout << endl
        //      << "-----------" << endl;

        //获得二阶导的向量
        vector<double> M = solve(A, b);
        // vec_out(M);
        // cout << endl
        //      << "-----------" << endl;

        //无结样条/周期样条
        if (type == notAknot || type == periodic)
        {
            //根据二阶导向量补充一阶导的信息
            // f[x_1, x_1] = f[x_1, x_2] - (2M_1 + M_2) / 6 * (x_2 - x_1)
            double f_prime = tableOfDividedDiffs_[1][1] - (2 * M[0] + M[1]) / 6 * (points[1] - points[0]);
            interpC_.interpSites[points[0]].push_back(f_prime);
        }

        //推导多项式
        for (int i = 0; i < N - 1; i++)
        {
            //将double转换为Vec1
            vector<double> coef = res_spline.pieceCalculate(interpC_, M, i).get_coefficients();
            vector<Vec1> _coef;
            for (int j = 0; j < coef.size(); j++)
            {
                _coef.push_back(coef[j]);
            }
            res_spline[i] = Polynomial<Ord, Vec1>(_coef);
        }

        break;
    }
    case cardinalB:
    {
        //三次B样条，在整数点插值
        if (Ord == 3 && type == complete)
        {
            //插值点的个数
            int N = interpC_.interpSites.size();
            vector<double> b(N);
            vector<vector<double>> M(N, vector<double>(N));

            //系数多项式
            for (int i = 1; i < N - 1; i++)
            {
                M[i][i - 1] = M[i][i + 1] = 1;
                M[i][i] = 4;
            }
            M[0][0] = M[N - 1][N - 1] = 4;
            M[0][1] = M[N - 1][N - 2] = 2;

            //目标向量
            for (int i = 1; i < N - 1; i++)
            {
                b[i] = 6 * interpC_.interpSites[points[i]][0];
            }
            b[0] = 6 * interpC_.interpSites[points[0]][0] + 2 * interpC_.interpSites[points[0]][1];
            b[N - 1] = 6 * interpC_.interpSites[points[N - 1]][0] - 2 * interpC_.interpSites[points[N - 1]][1];

            //获得样条系数
            vector<double> a(N + 2);
            vector<double> tmp_a = solve(M, b);
            for (int i = 0; i < N; i++)
            {
                a[i + 1] = tmp_a[i];
            }
            a[0] = a[2] - 2 * interpC_.interpSites[points[0]][1];
            a[N + 1] = a[N - 1] + 2 * interpC_.interpSites[points[N - 1]][1];

            //累计样条函数
            for (int i = -1; i < N + 1; i++)
            {
                res_spline = res_spline + B_spline<3>(i) * a[i + 1];
            }
            //根据左端点的位置平移
            res_spline = (res_spline >> (points[0] - 1));
        }
        //二次B样条，在整数区间中点插值
        if (Ord == 2 && type == notAknot)
        {
            //注意一共有N个插值点，但两端都是整数点，需要排除
            //插值点的个数
            int N = interpC_.interpSites.size() - 1;
            vector<double> b(N - 1);
            vector<vector<double>> M(N - 1, vector<double>(N - 1));

            //系数多项式
            for (int i = 1; i < N - 2; i++)
            {
                M[i][i - 1] = M[i][i + 1] = 1;
                M[i][i] = 6;
            }
            M[0][0] = M[N - 2][N - 2] = 5;
            M[0][1] = M[N - 2][N - 3] = 1;

            //目标向量
            for (int i = 1; i < N - 2; i++)
            {
                b[i] = 8 * interpC_.interpSites[points[i + 1]][0];
            }
            b[0] = 8 * interpC_.interpSites[points[1]][0] - 2 * interpC_.interpSites[points[0]][0];
            b[N - 2] = 8 * interpC_.interpSites[points[N - 1]][0] - 2 * interpC_.interpSites[points[N]][0];

            //获得样条系数
            vector<double> a(N + 1);
            vector<double> tmp_a = solve(M, b);
            for (int i = 0; i < N - 1; i++)
            {
                a[i + 1] = tmp_a[i];
            }
            a[0] = 2 * interpC_.interpSites[points[0]][0] - a[1];
            a[N] = 2 * interpC_.interpSites[points[N]][0] - a[N - 1];

            //累计样条函数
            for (int i = 0; i < N + 1; i++)
            {
                res_spline = res_spline + B_spline<2>(i) * a[i];
            }
            //根据左端点的位置平移
            res_spline = (res_spline >> (points[0] - 1));
        }
        break;
    }
    }

    return res_spline;
}

/**
 * @date   2021/11/23 23:32:10
 *
 * @arg    const vector<Vec<double, 2>> &points, BCType type
 * @return Spline<2, Ord, ppForm>
 *
 */

//曲线拟合
template <int Ord>
Spline<2, Ord, ppForm> fitCurve(const vector<Vec<double, 2>> &points, BCType type)
{
    vector<double> lengths(points.size());

    //计算累计弦长
    lengths[0] = 0;
    for (int i = 1; i < points.size(); i++)
    {
        double dx = points[i][0] - points[i - 1][0];
        double dy = points[i][1] - points[i - 1][1];
        lengths[i] = lengths[i - 1] + sqrt(dx * dx + dy * dy);
    }

    //分别计算插值多项式
    INTERPC interpC_x;
    INTERPC interpC_y;
    interpC_x.points = lengths;
    interpC_y.points = lengths;
    for (int i = 0; i < lengths.size(); i++)
    {
        interpC_x.interpSites[lengths[i]] = {points[i][0]};
        interpC_y.interpSites[lengths[i]] = {points[i][1]};
    }

    Spline<1, Ord, ppForm> res_spline1 = interpolate<Ord, ppForm>(interpC_x, type);
    Spline<1, Ord, ppForm> res_spline2 = interpolate<Ord, ppForm>(interpC_y, type);

    vector<Polynomial<Ord, Vec2>> spline_poly;
    for (int i = 0; i < res_spline1.size(); i++)
    {
        //合并多项式为二维多项式
        spline_poly.push_back(overlying<Ord, 1, 1>(res_spline1.spline_poly[i], res_spline2.spline_poly[i]));
    }
    Spline<2, Ord, ppForm> res_spline(spline_poly);
    res_spline.interval = lengths;
    return res_spline;
}

void vec_out(vector<double> &vec)
{
    cout << "(";
    for (int i = 0; i < vec.size(); i++)
    {
        cout << vec[i] << ",";
    }
    cout << "\b)";
}

void mat_out(vector<vector<double>> &mat)
{
    for (int i = 0; i < mat.size(); i++)
    {
        vec_out(mat[i]);
        cout << endl;
    }
}

double max_norm(vector<double> &vec)
{
    double m = 0;
    for (int i = 0; i < vec.size(); i++)
    {
        if (fabs(vec[i]) > m)
        {
            m = vec[i];
        }
    }
    return m;
}

vector<double> solve(vector<vector<double>> &_A, vector<double> &_b)
{
    int size = _b.size();

    double A[size * size];
    double B[size];

    for (int i = 0; i < size; i++)
    {
        for (int j = 0; j < size; j++)
        {
            A[i * size + j] = _A[i][j];
        }
        B[i] = _b[i];
    }

    int ipiv[size];
    int N = size;
    int NRHS = 1;
    int LDA = size;
    int LDB = 1;
    int INFO = 0;

    LAPACKE_dgesv(LAPACK_ROW_MAJOR, N, NRHS, A, LDA, ipiv, B, LDB);

    for (int i = 0; i < size; i++)
    {
        _b[i] = B[i];
    }

    return _b;
}

#endif
